#include "fd_bits.h"

ulong
fd_ulong_approx_sqrt( ulong x ) {

  /* Theory:

     Break a non-zero x exactly into:

       x = 2^e ( 1 + m/2^63 )

     where e is 6-bit and m is 63-bit.  Further, break e exactly into:

       e = 2 q + r

     where q is 5-bit and r in 1-bit .  Then, break m exactly into:

       m = u 2^61 + h 2^29 + l

     where u is 2-bit, h is 32-bit and l is 29-bit.  Then:

       y = sqrt( x )
         = 2^q 2^(r/2) sqrt( 1 + u/4 + h/2^34 + l/2^63 )

     We ignore the l term as very small.  For given r and u, y is
     roughly linearly with respect to h.  Let i be a 3-bit table index
     such that:

       i = 4 r + u

     and approximate:

       y ~ 2^q (alpha_i + h beta_i / 2^32)

     where:

       alpha_i = 2^(r/2) sqrt( 1 + u/4 )
       beta_i  = alpha_{i+1} - alpha_i

     The right hand term is in [1,2] for h in [0,2^32) after rounding.
     To compute this in 64-bit integer math then, we scale the right
     hand term by 2^62 yielding:

       y ~ round( 2^62 (alpha_i + beta_i h) / 2^(62-q) )
         ~ floor( (a_i + b_i h) / 2^(62-q) + 0.5 )
         ~ (a_i + b_i h + 2^(61-q)) / 2^(62-q)

     where:

       a_i = round( 2^62 alpha_i )
       b_i = round( 2^30  beta_i ) */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  int   e = fd_ulong_find_msb( x );         /* In [0,63]      */
  ulong m = x << (63-e);                    /* In [2^63,2^64) */
  int   q = e >> 1;                         /* In [0,31]      */
  int   r = e & 1;                          /* In [0,1]       */
  int   i = (4*r) + (int)((m >> 61) & 3UL); /* In [0,8)       */
  ulong h = (ulong)(uint)(m >> 29);         /* In [0,2^32)    */

  /* These tables are autogenerated using long double arithmetic */

  static ulong const a[8] = { 4611686018427387904UL, 5156021714044493574UL, 5648138799537240564UL, 6100687164736251614UL,
                              6521908912666391106UL, 7291715835891894856UL, 7987674492471257551UL, 8627674528165471361UL };
  static ulong const b[8] = {           126738030UL,           114579938UL,           105367127UL,            98073331UL,
                                        179234641UL,           162040502UL,           149011620UL,           138696634UL };

  return (a[i] + b[i]*h + (1UL<<(61-q))) >> (62-q);
}

ulong
fd_ulong_round_sqrt( ulong x ) {

  /* Theory:

     To find the integer y such that y = round( sqrt(x) ) with ties
     rounding toward zero, note:

          (y-0.5)^2 < x <= (y+0.5)^2
       -> y^2 - y + 0.25 < x <= y^2 + y + 0.25

     Since x and y are integral, ties can never occur.  Thus:

       -> y^2 - y + 0.25 < x < y^2 + y + 0.25

     Similarly, integral x and y means we can eliminate the 0.25 adds:

       -> y^2 - y + 1 <= x < y^2 + y + 1

     Let y^2 - y + 1 = x - r.

       -> x - r <= x < x - r + 2 y
       -> 0 <= r < 2 y

     That is, if y is round( sqrt(x) ) we must have:

        y^2 - y + 1 = x - r for some r in [0,2 y).

     Suppose we have a reasonable initial guess.  Applying
     Newton-Raphson to:

       f(y) = y^2 - x

     for real valued x and y yields:

       y1 = y0 + (x - y0^2)/(2 y0)

     Naively generalizing to integral x and y yields:

       y1 = y0 + floor( (x - y0^2)/(2 y0) )

     If this converges, at convergence we have:

          (x - y^2 - r) / (2 y) = 0 for some r in [0,2 y)
       -> y^2 = x - r for some r in [0,2 y)

     This the naive implementation doesn't converge to the desired y.
     But if we tweak the N-R iteration to:

       y1 = y0 + floor( (x - y0^2 + y0 - 1) / (2 y0) )

     At convergence we have:

          x - y^2 + y - 1 - r = 0 for some r in [0,2 y)
       -> y^2 - y + 1 = x - r     for some r in [0,2 y).

     which is the desired solution.

     N-R converges quadratically (i.e. the number of correct bits
     roughly doubles each iteration given a reasonable initial guess).
     The initial guess here is accurate to 4+ bits and the solution
     requires ~32 bits accuracy.  So this will complete with at most ~3
     divisions.  Though it is slightly faster for large x to use an
     unrolled implementation (e.g. iterate 3 quickly), we use a variable
     iteration count to be absolutely sure of convergence and to be
     faster when fed smaller inputs (as is done in typical usage ...
     where O(1) call is made with a moderate sized x).

     Since our initial guess is close, y^2 fits within 65 bits.  Thus,
     the subtraction to compute x - y^2 will "do the right thing" with
     no special handling needed for wrapping.

     Since modern C/C++ use round toward zero for integer division, we
     have to tweak the numerator when it is negative to convert to floor
     style division. */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  ulong y = fd_ulong_approx_sqrt( x );

  for(;;) {
    long den = (long)(y << 1);
    long num = (long)(x - y*y + y - 1UL);
    if( ((0L<=num) & (num<den)) ) break; /* Branch prob depends on x's distribution */
    y += (ulong)((num - fd_long_if( num>=0L, 0L, den-1L )) / den);
  }

  return y;
}

ulong
fd_ulong_floor_sqrt( ulong x ) {

  /* Theory:

     To find the integer y such that y = floor( sqrt(x) ), note:

          y^2 <= x < (y+1)^2
       -> y^2 <= x < y^2 + 2 y + 1

     Let y^2 = x - r.

       -> x - r <= x < x - r + 2 y + 1
       -> 0 <= r < 2 y + 1

     That is, if y is floor( sqrt(x) ) we must have:

        y^2 = x - r for some r in [0,2 y + 1).

     Like the round implementation, naively generalizing N-R
     to integer doesn't converge to the correct y.  But if
     we tweak the N-R iteration to:

       y1 = y0 + floor( (x - y0^2) / (2 y0 + 1) )

     At convergence we have:

          x - y^2 - r = 0 for some r in [0,2 y+1)
       -> y^2 = x - r     for some r in [0,2 y).

     which is the desired solution.

     Similar considerations for convergence, wrapping and trunc vs floor
     division as the above. */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  ulong y = fd_ulong_approx_sqrt( x );

  for(;;) {
    long den = (long)((y << 1) + 1UL);
    long num = (long)(x - y*y);
    if( ((0L<=num) & (num<den)) ) break; /* Branch prob depends on x's distribution */
    y += (ulong)((num - fd_long_if( num>=0L, 0L, den-1L )) / den);
  }

  return y;
}

ulong
fd_ulong_ceil_sqrt( ulong x ) {

  /* Theory:

     To find the integer y such that y = ceil( sqrt(x) ), note that
     positive x:

          (y-1)^2 < x <= y
       -> y^2 - 2 y + 1 < x <= y^2

     Let y = x + r.

       -> x + r - 2 y + 1 < x <= x + r
       -> -2 y + 1 < -r <= 0
       -> 0 <= r < 2 y - 1

     That is, if y is ceil( sqrt(x) ) we must have:

        y^2 = x + r for some r in [0,2 y - 1).

     Like the round implementation, naively generalizing N-R
     to integer doesn't converge to the correct y.  But if
     we tweak the N-R iteration to:

       y1 = y0 + ceil ( (x - y0^2           ) / (2 y0 - 1) )
          = y0 + floor( (x - y0^2 + 2 y0 - 2) / (2 y0 - 1) )

     At convergence we have:

          x - y^2 + r = 0 for some r in [0,2 y-1)
       -> y^2 = x + r     for some r in [0,2 y-1).

     which is the desired solution.

     Similar considerations for convergence, wrapping and trunc vs floor
     division as the above. */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  ulong y = fd_ulong_approx_sqrt( x );

  for(;;) {
    ulong tmp = y << 1;
    long  den = (long)(tmp - 1UL);
    long  num = (long)(x - y*y + tmp - 2UL);
    if( ((0L<=num) & (num<den)) ) break; /* Branch prob depends on x's distribution */
    y += (ulong)((num - fd_long_if( num>=0L, 0L, den-1L )) / den);
  }

  return y;
}

ulong
fd_ulong_approx_cbrt( ulong x ) {

  /* This is basically the same as the approx_sqrt but only uses 1
     significant bit for u. */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  int   e = fd_ulong_find_msb( x );       /* In [0,63]      */
  ulong m = x << (63-e);                  /* In [2^63,2^64) */
  int   q = (int)(((uint)e) / 3U);        /* In [0,21]      */
  int   r = (int)(((uint)e) % 3U);        /* In [0,2)       */
  int   i = 2*r + (int)((m >> 62) & 1UL); /* In [0,6)       */
  ulong h = (ulong)(uint)(m >> 30);       /* In [0,2^32)    */

  static ulong const a[6] = {  4611686018427387904UL,  5279062667477898215UL,  5810360290122541960UL,
                               6651202178469583219UL,  7320595236998672907UL,  8379989631760464848UL };
  static ulong const b[6] = {            155385734UL,            123702367UL,            195773758UL,
                                         155855216UL,            246659478UL,            196365268UL };

  return (a[i] + b[i]*h + (1UL<<(61-q))) >> (62-q);
}

ulong
fd_ulong_round_cbrt( ulong x ) {

  /* This works the same as the above with:
          (y-0.5)^3 <= x < (y+0.5)^3
       -> 4 y^3 - 6 y^2 + 3 y = 4 x - r for some r in [0,12 y^2 + 1)
     and Newton-Raphson generalized to integer arithmetic applied to:
        f(y) = y^3 - x
     yielding:
        y1 = y0 + floor( 4 x - 4 y0^3 + 6 y0^2 - 3 y0) / (12 y0^2 + 1) ) */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  ulong y = fd_ulong_approx_cbrt( x );

  for(;;) {
    ulong ysq = y*y;
    long  num = (long)(4UL*(x - y*ysq) + 6UL*ysq - 3UL*y);
    long  den = (long)(12UL*ysq + 1UL);
    if( ((0L<=num) & (num<den)) ) break; /* Branch prob depends on x's distribution */
    y += (ulong)((num - fd_long_if( num>=0L, 0L, den-1L )) / den);
  }

  return y;
}

ulong
fd_ulong_floor_cbrt( ulong x ) {

  /* This works the same as the above with:
          y^3 <= x < (y+1)^3
       -> y^3 = x - r for some r in [0,3 y^2 + 3 y + 1)
     and Newton-Raphson generalized to integer arithmetic applied to:
        f(y) = y^3 - x
     yielding:
        y1 = y0 + floor( (x - y^3) / (3 y^2 + 3 y + 1) ) */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  ulong y = fd_ulong_approx_cbrt( x );

  for(;;) {
    ulong ysq = y*y;
    long  num = (long)(x - y*ysq);
    long  den = (long)(3UL*(ysq+y) + 1UL);
    if( ((0L<=num) & (num<den)) ) break; /* Branch prob depends on x's distribution */
    y += (ulong)((num - fd_long_if( num>=0L, 0L, den-1L )) / den);
  }

  return y;
}

ulong
fd_ulong_ceil_cbrt( ulong x ) {

  /* This works the same as the above with:
          (y-1)^3 < x <= y^3
       -> y^3 = x + r for some r in [0,3 y^2 - 3 y + 1)
     and Newton-Raphson generalized to integer arithmetic applied to:
        f(y) = y^3 - x
     yielding:
        y1 = y0 + ceil ( (x - y^3              ) / (3 y^2 - 3 y + 1) )
           = y0 + floor( (x - y^3 + 3 y^2 - 3 y) / (3 y^2 - 3 y + 1) ) */

  if( FD_UNLIKELY( !x ) ) return 0UL;

  ulong y = fd_ulong_approx_cbrt( x );

  for(;;) {
    ulong ysq = y*y;
    ulong tmp = 3UL*(ysq-y);
    long  num = (long)(x - y*ysq + tmp);
    long  den = (long)(tmp + 1UL);
    if( ((0L<=num) & (num<den)) ) break; /* Branch prob depends on x's distribution */
    y += (ulong)((num - fd_long_if( num>=0L, 0L, den-1L )) / den);
  }

  return y;
}
